#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LangGraph多智能体HR系统 - 修复版
解决无限循环问题，简化架构，提升稳定性
"""

import os
import json
import logging
import asyncio
from datetime import datetime
from typing import Dict, List, Any, Literal, TypedDict, Annotated
from typing_extensions import NotRequired

# LangGraph核心模块
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

# LangChain核心模块
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

# 数据管理
from data_manager import data_manager
from agent_prompts import agent_prompts

# 日志系统
from logger_config import hr_logger, Stage, performance_monitor
from timeout_monitor import timeout_monitor, TimeoutType, TimeoutContext

# 配置基础日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# ==================== 1. 系统状态定义 ====================

class MultiAgentState(TypedDict):
    """多智能体系统全局状态"""
    # 基础会话信息
    session_id: str
    user_id: str
    timestamp: str
    
    # 用户交互
    original_query: str
    messages: Annotated[List[Any], add_messages]
    
    # 业务数据
    candidates: NotRequired[List[Dict[str, Any]]]
    analysis_results: NotRequired[Dict[str, Any]]
    ui_commands: NotRequired[List[Dict[str, Any]]]
    
    # 系统状态
    next_agent: NotRequired[str]
    final_response: NotRequired[str]
    status: str
    processing_stage: NotRequired[str]  # 新增：处理阶段标记

# ==================== 2. 工具定义 ====================

@tool
def search_candidates_by_criteria(skills: list = None, min_experience: int = 0, location: str = "", limit: int = 10) -> List[Dict[str, Any]]:
    """根据条件搜索候选人
    
    Args:
        skills: 技能列表，如 ["Python", "Java"]
        min_experience: 最少工作经验年数，默认0年
        location: 地点要求，如 "北京"
        limit: 限制返回数量，默认10个
    
    Returns:
        匹配的候选人列表
    """
    if skills is None:
        skills = []
    
    criteria = {
        'skills': skills,
        'min_experience': min_experience,
        'location': location
    }
    
    hr_logger.tool_execution("search_candidates_by_criteria", criteria)
    
    matching_candidates = []
    
    for resume in data_manager.resumes:
        # 技能匹配
        skill_match = False
        if not skills:
            skill_match = True
        else:
            resume_skills_str = ' '.join(resume.skills) if isinstance(resume.skills, list) else str(resume.skills)
            for skill in skills:
                if skill.lower() in resume_skills_str.lower():
                    skill_match = True
                    break
        
        # 经验匹配
        experience_years = resume.experience_years or 0
        experience_match = experience_years >= min_experience
        
        # 地点匹配  
        location_match = True
        if location:
            location_match = location.lower() in resume.location.lower()
        
        if skill_match and experience_match and location_match:
            # 计算匹配分数
            score = 70  # 基础分数
            if skills:
                # 根据技能匹配度加分
                matched_skills = sum(1 for skill in skills if skill.lower() in resume_skills_str.lower())
                score += (matched_skills / len(skills)) * 20
            if experience_years >= min_experience:
                score += 10
            
            matching_candidates.append({
                'id': resume.id,
                'name': resume.name or f"候选人_{resume.id[:8]}",
                'skills': resume.skills if isinstance(resume.skills, list) else [str(resume.skills)],
                'experience': experience_years,
                'location': resume.location or "未知",
                'education': getattr(resume, 'education', getattr(resume, 'qualification', '未知')),
                'match_score': min(95, score),
                'resume_preview': resume.content[:200] + "..." if len(resume.content) > 200 else resume.content
            })
    
    # 按匹配分数排序
    matching_candidates.sort(key=lambda x: x['match_score'], reverse=True)
    
    logger.info(f"✅ 找到 {len(matching_candidates)} 个匹配候选人，返回前 {limit} 个")
    return matching_candidates[:limit]

@tool
def generate_interview_questions(candidate_id: str, position: str = "软件工程师", difficulty: str = "中等") -> Dict[str, Any]:
    """为候选人生成面试题目
    
    Args:
        candidate_id: 候选人ID
        position: 岗位名称，如 "Python工程师", "前端工程师"  
        difficulty: 难度级别，"简单", "中等", "困难"
    
    Returns:
        生成的面试题目列表
    """
    hr_logger.tool_execution("generate_interview_questions", {
        "candidate_id": candidate_id,
        "position": position,
        "difficulty": difficulty
    })
    
    # 根据候选人技能生成针对性面试题
    candidate = None
    for resume in data_manager.resumes:
        if resume.id == candidate_id:
            candidate = resume
            break
    
    if not candidate:
        return {"error": f"候选人 {candidate_id} 不存在"}
    
    # 基于候选人技能生成题目
    questions = []
    skills = candidate.skills if isinstance(candidate.skills, list) else []
    
    # Python相关题目
    if any('python' in str(skill).lower() for skill in skills):
        if difficulty == "简单":
            questions.extend([
                "请解释Python中list和tuple的区别",
                "什么是Python的装饰器？请举个例子"
            ])
        elif difficulty == "中等":
            questions.extend([
                "请解释Python的GIL（全局解释器锁）",
                "如何实现Python中的单例模式？"
            ])
        else:  # 困难
            questions.extend([
                "请解释Python的元类（metaclass）概念",
                "如何优化Python程序的性能？"
            ])
    
    # JavaScript相关题目
    if any('javascript' in str(skill).lower() for skill in skills):
        questions.extend([
            "解释JavaScript中的闭包概念",
            "什么是JavaScript的事件循环？"
        ])
    
    # 通用编程题目
    questions.extend([
        "设计一个简单的用户管理系统的数据库表结构",
        "描述你在项目中遇到的最大技术挑战及解决方案"
    ])
    
    return {
        'candidate_id': candidate_id,
        'candidate_name': candidate.name or f"候选人_{candidate_id[:8]}",
        'position': position,
        'difficulty': difficulty,
        'questions': questions[:5],  # 限制题目数量
        'generated_time': datetime.now().isoformat(),
        'skills_based': [skill for skill in skills if skill.lower() in ['python', 'javascript', 'java']]
    }

@tool
def update_candidate_record(candidate_id: str, status: str = "", notes: str = "", tags: str = "") -> Dict[str, Any]:
    """更新候选人记录
    
    Args:
        candidate_id: 候选人ID
        status: 状态更新，如 "面试中", "已录用", "不合适"
        notes: 备注信息
        tags: 标签，多个标签用逗号分隔
    
    Returns:
        更新操作结果
    """
    updates = {}
    if status:
        updates['status'] = status
    if notes:
        updates['notes'] = notes
    if tags:
        updates['tags'] = tags.split(',')
    
    hr_logger.tool_execution("update_candidate_record", {
        "candidate_id": candidate_id,
        "updates": updates
    })
    
    operation_log = {
        'candidate_id': candidate_id,
        'timestamp': datetime.now().isoformat(),
        'updates': updates,
        'status': 'success'
    }
    
    hr_logger.info(f"候选人记录更新成功: {operation_log}")
    return operation_log

# 工具列表
hr_tools = [
    search_candidates_by_criteria,
    generate_interview_questions,
    update_candidate_record
]

# ==================== 3. 大模型配置 ====================

def create_multi_agent_model():
    """创建多智能体模型"""
    hr_logger.info("开始创建多智能体模型")
    
    # 检查API密钥
    dashscope_key = os.getenv("DASHSCOPE_API_KEY")
    
    if dashscope_key:
        hr_logger.info("使用阿里云通义千问模型", {"model": "qwen-plus", "api_base": "dashscope"})
        
        with TimeoutContext(timeout_monitor, TimeoutType.MODEL_INVOKE, "create_model_dashscope"):
            model = ChatOpenAI(
                model="qwen-plus",
                openai_api_key=dashscope_key,
                openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
                temperature=0.1,
                timeout=30,
                max_retries=2
            )
    else:
        hr_logger.error("未找到DASHSCOPE_API_KEY环境变量")
        raise Exception("未找到API密钥，请设置 DASHSCOPE_API_KEY 环境变量")
    
    # 绑定工具
    model_with_tools = model.bind_tools(hr_tools)
    hr_logger.info("模型创建完成", {"tools_count": len(hr_tools)})
    
    return model_with_tools

# 全局模型实例
multi_agent_model = create_multi_agent_model()

# ==================== 4. 智能体节点定义 ====================

@performance_monitor("interaction_module")
def interaction_module(state: MultiAgentState) -> Dict[str, Any]:
    """交互模块 - 处理用户输入，启动系统"""
    hr_logger.stage_start(Stage.INTERACTION, "开始处理用户输入")
    
    user_query = state["original_query"]
    session_id = state.get("session_id", f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
    user_id = state.get("user_id", "hr_user_001")
    
    hr_logger.user_interaction(user_id, "提交查询", {
        "query_length": len(user_query),
        "query_preview": user_query[:100]
    })
    
    hr_logger.stage_end(Stage.INTERACTION, "用户输入处理完成，转交给Orchestrator")
    
    return {
        "session_id": session_id,
        "messages": [HumanMessage(content=user_query)],
        "next_agent": "orchestrator",
        "status": "processing",
        "processing_stage": "interaction_completed"
    }

@performance_monitor("orchestrator_agent")
def orchestrator_agent(state: MultiAgentState) -> Dict[str, Any]:
    """Orchestrator Agent - 总调度中心（简化版）"""
    hr_logger.stage_start(Stage.ORCHESTRATOR, "开始任务分析和计划制定")
    
    # 简化的提示词，专注于工具调用
    orchestrator_prompt = """你是专业的HR智能助手，负责理解用户需求并调用相应的工具。

对于用户的请求，你需要：
1. 分析用户意图
2. 调用合适的工具来完成任务
3. 不要只用文字回复，必须调用工具

可用工具：
- search_candidates_by_criteria: 搜索和筛选候选人
- generate_interview_questions: 生成面试题
- update_candidate_record: 更新候选人状态

示例：
- "筛选python工程师" → 调用 search_candidates_by_criteria(skills=["Python"])
- "给候选人12345生成面试题" → 调用 generate_interview_questions(candidate_id="12345")
- "将候选人设为面试中" → 调用 update_candidate_record(status="面试中")

请根据用户需求调用合适的工具。"""

    messages = [SystemMessage(content=orchestrator_prompt)] + state["messages"]
    
    try:
        response = multi_agent_model.invoke(messages)
        
        if hasattr(response, 'tool_calls') and response.tool_calls:
            # 需要调用工具
            logger.info(f"🔧 Orchestrator决定调用工具: {[tc['name'] for tc in response.tool_calls]}")
            
            return {
                "messages": [response],
                "next_agent": "tool_executor",
                "processing_stage": "tool_execution_needed"
            }
        else:
            # 直接回复，进入完成阶段
            return {
                "messages": [response],
                "next_agent": "completion",
                "processing_stage": "direct_response"
            }
    
    except Exception as e:
        hr_logger.error("Orchestrator大模型调用失败", exception=e)
        return {
            "messages": [AIMessage(content=f"处理失败：{str(e)}")],
            "next_agent": "completion",
            "status": "error"
        }

def tool_executor(state: MultiAgentState) -> Dict[str, Any]:
    """工具执行器（修复版）"""
    logger.info("🛠️ Tool Executor: 执行工具调用")
    
    last_message = state["messages"][-1]
    
    if not hasattr(last_message, 'tool_calls') or not last_message.tool_calls:
        logger.error("❌ 没有找到工具调用请求")
        return {
            "messages": [AIMessage(content="工具执行失败：没有找到工具调用请求")],
            "next_agent": "completion"
        }
    
    tool_results = []
    candidates_found = []
    interview_questions = None
    
    for tool_call in last_message.tool_calls:
        tool_name = tool_call["name"]
        tool_args = tool_call["args"]
        tool_id = tool_call.get("id", "unknown")
        
        logger.info(f"🔧 执行工具: {tool_name} 参数: {tool_args}")
        
        # 查找对应的工具函数
        tool_function = None
        for tool in hr_tools:
            if tool.name == tool_name:
                tool_function = tool
                break
        
        if tool_function:
            result = tool_function.invoke(tool_args)
            
            # 处理结果
            if tool_name == "search_candidates_by_criteria" and isinstance(result, list):
                candidates_found = result
            elif tool_name == "generate_interview_questions" and isinstance(result, dict):
                interview_questions = result
            
            # 创建工具消息
            from langchain_core.messages import ToolMessage
            tool_message = ToolMessage(
                content=str(result),
                tool_call_id=tool_id,
                name=tool_name
            )
            tool_results.append(tool_message)
            
            logger.info(f"✅ 工具 {tool_name} 执行成功")
    
    # 构建最终响应
    if candidates_found:
        response_text = f"找到 {len(candidates_found)} 个匹配的候选人：\n\n"
        for i, candidate in enumerate(candidates_found[:5], 1):
            response_text += f"{i}. {candidate['name']} (匹配度: {candidate['match_score']}%)\n"
            response_text += f"   技能: {', '.join(candidate['skills'][:3])}\n"
            response_text += f"   经验: {candidate['experience']}年\n\n"
    elif interview_questions:
        response_text = f"为候选人 {interview_questions['candidate_name']} 生成的面试题：\n\n"
        for i, q in enumerate(interview_questions['questions'], 1):
            response_text += f"{i}. {q}\n"
    else:
        response_text = "操作已完成。"
    
    # 构建UI指令
    ui_commands = []
    if candidates_found:
        ui_commands.append({
            "command": "UPDATE_CANDIDATE_LIST",
            "payload": candidates_found
        })
    if interview_questions:
        ui_commands.append({
            "command": "SHOW_INTERVIEW_QUESTIONS",
            "payload": interview_questions
        })
    
    return {
        "messages": tool_results + [AIMessage(content=response_text)],
        "candidates": candidates_found,
        "analysis_results": {"interview_questions": interview_questions} if interview_questions else {},
        "ui_commands": ui_commands,
        "next_agent": "completion",  # 固定路由到completion，避免循环
        "processing_stage": "tool_execution_completed"
    }

def completion_handler(state: MultiAgentState) -> Dict[str, Any]:
    """完成处理器 - 生成最终响应"""
    logger.info("🏁 Completion Handler: 生成最终响应")
    
    # 提取最终的AI消息作为响应
    final_message = ""
    for msg in reversed(state["messages"]):
        if isinstance(msg, AIMessage) and msg.content:
            # 跳过包含工具调用的消息
            if not (hasattr(msg, 'tool_calls') and msg.tool_calls):
                final_message = msg.content
                break
    
    if not final_message:
        final_message = "任务已完成，感谢使用HR智能助手！"
    
    return {
        "final_response": final_message,
        "status": "completed",
        "next_agent": END,
        "processing_stage": "completed"
    }

# ==================== 5. 路由函数（简化版）====================

def route_next_agent(state: MultiAgentState) -> str:
    """智能体路由函数（修复版）"""
    next_agent = state.get("next_agent", "completion")
    logger.info(f"🗺️ 路由到下一个智能体: {next_agent}")
    
    # 简化路由逻辑，避免循环
    valid_routes = ["orchestrator", "tool_executor", "completion"]
    if next_agent in valid_routes:
        return next_agent
    else:
        logger.warning(f"无效路由目标: {next_agent}, 默认路由到completion")
        return "completion"

# ==================== 6. 多智能体系统构建（简化版）====================

class MultiAgentHRSystem:
    """多智能体HR系统（修复版）"""
    
    def __init__(self):
        self.graph = None
        self.initialized = False
    
    def initialize(self) -> bool:
        """初始化多智能体系统"""
        logger.info("🚀 初始化多智能体HR系统...")
        
        # 初始化数据管理器
        data_manager.initialize()
        logger.info(f"📊 数据加载完成: {len(data_manager.resumes)}个简历, {len(data_manager.jobs)}个岗位")
        
        # 构建简化的LangGraph
        workflow = StateGraph(MultiAgentState)
        
        # 添加智能体节点
        workflow.add_node("interaction", interaction_module)
        workflow.add_node("orchestrator", orchestrator_agent)
        workflow.add_node("tool_executor", tool_executor)
        workflow.add_node("completion", completion_handler)
        
        # 设置入口点
        workflow.set_entry_point("interaction")
        
        # 简化的路由结构，避免循环
        workflow.add_conditional_edges(
            "interaction",
            route_next_agent,
            {
                "orchestrator": "orchestrator",
                "completion": "completion"
            }
        )
        
        workflow.add_conditional_edges(
            "orchestrator", 
            route_next_agent,
            {
                "tool_executor": "tool_executor",
                "completion": "completion"
            }
        )
        
        workflow.add_conditional_edges(
            "tool_executor",
            route_next_agent,
            {
                "completion": "completion"  # 只能路由到completion，避免循环
            }
        )
        
        # 编译图
        self.graph = workflow.compile()
        self.initialized = True
        
        logger.info("✅ 多智能体HR系统初始化完成！")
        logger.info("🎯 简化架构: Interaction → Orchestrator → Tools → Completion")
        return True
    
    async def process_request(self, user_query: str, user_id: str = "hr_user_001") -> Dict[str, Any]:
        """处理用户请求"""
        if not self.initialized:
            return {"success": False, "error": "系统未初始化"}
        
        logger.info(f"🎯 开始处理用户请求: {user_query}")
        
        # 初始状态
        initial_state = {
            "original_query": user_query,
            "user_id": user_id,
            "timestamp": datetime.now().isoformat(),
            "status": "started",
            "messages": []
        }
        
        # 执行多智能体工作流 - 降低递归限制
        config = {
            "recursion_limit": 10,  # 降低递归限制
            "max_execution_time": 60  # 降低超时时间
        }
        final_state = await self.graph.ainvoke(initial_state, config=config)
        
        # 构建响应
        response = {
            "success": True,
            "message": final_state.get("final_response", "任务完成"),
            "session_id": final_state.get("session_id"),
            "status": final_state.get("status", "completed"),
            "ui_commands": final_state.get("ui_commands", []),
            "analysis_results": final_state.get("analysis_results", {}),
            "candidates": final_state.get("candidates", []),
            "agent_workflow": "Simplified Multi-Agent: Interaction → Orchestrator → Tools → Completion"
        }
        
        logger.info("✅ 多智能体请求处理完成")
        return response

# 全局系统实例
multi_agent_hr_system = MultiAgentHRSystem()

if __name__ == "__main__":
    import asyncio
    
    async def main():
        # 初始化系统
        multi_agent_hr_system.initialize()
        
        # 测试请求
        test_query = "筛选出会Python的工程师"
        result = await multi_agent_hr_system.process_request(test_query)
        
        print("🎯 测试结果:")
        print(json.dumps(result, indent=2, ensure_ascii=False))
    
    asyncio.run(main())





